"""
缓存与断点续跑模块
用于缓存各个阶段的产物，支持断点续跑和增量更新
"""

import hashlib
import json
import os
from pathlib import Path
from typing import Any, Dict, Optional, List
from datetime import datetime
import shutil


class CacheManager:
    """缓存管理器"""

    CACHE_VERSION = "1.0"

    def __init__(self, cache_dir: Path):
        """
        初始化缓存管理器

        Args:
            cache_dir: 缓存目录
        """
        self.cache_dir = cache_dir
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.metadata_file = cache_dir / "cache_metadata.json"
        self.metadata = self._load_metadata()

    def _load_metadata(self) -> Dict[str, Any]:
        """加载缓存元数据"""
        if self.metadata_file.exists():
            try:
                with open(self.metadata_file, "r", encoding="utf-8") as f:
                    return json.load(f)
            except Exception:
                pass
        return {
            "version": self.CACHE_VERSION,
            "items": {},
            "created_at": datetime.now().isoformat(),
        }

    def _save_metadata(self) -> None:
        """保存缓存元数据"""
        with open(self.metadata_file, "w", encoding="utf-8") as f:
            json.dump(self.metadata, f, ensure_ascii=False, indent=2)

    def _get_file_hash(self, file_path: Path) -> str:
        """计算文件哈希值"""
        hasher = hashlib.sha256()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hasher.update(chunk)
        return hasher.hexdigest()

    def _get_content_hash(self, content: str) -> str:
        """计算内容哈希值"""
        return hashlib.sha256(content.encode("utf-8")).hexdigest()

    def _get_cache_key(self, key: str, params: Dict[str, Any]) -> str:
        """
        生成缓存键

        Args:
            key: 缓存键名
            params: 参数

        Returns:
            缓存键
        """
        # 将参数转换为字符串并计算哈希
        param_str = json.dumps(params, sort_keys=True, ensure_ascii=False)
        param_hash = hashlib.sha256(param_str.encode("utf-8")).hexdigest()[:16]
        return f"{key}_{param_hash}"

    def _get_cache_path(self, cache_key: str) -> Path:
        """获取缓存文件路径"""
        return self.cache_dir / f"{cache_key}.cache"

    def get(
        self,
        key: str,
        params: Dict[str, Any],
        output_path: Optional[Path] = None,
    ) -> Optional[Path]:
        """
        获取缓存

        Args:
            key: 缓存键名
            params: 参数
            output_path: 输出文件路径（如果指定，将从缓存复制到该路径）

        Returns:
            缓存文件路径（如果存在且有效）
        """
        cache_key = self._get_cache_key(key, params)
        cache_path = self._get_cache_path(cache_key)

        if not cache_path.exists():
            return None

        # 检查元数据
        item_meta = self.metadata["items"].get(cache_key)
        if not item_meta:
            return None

        # 检查文件是否还存在
        if not cache_path.exists():
            return None

        # 复制缓存文件到输出路径
        if output_path:
            output_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(cache_path, output_path)
            return output_path

        return cache_path

    def set(
        self,
        key: str,
        params: Dict[str, Any],
        file_path: Optional[Path] = None,
        content: Optional[str] = None,
    ) -> bool:
        """
        设置缓存

        Args:
            key: 缓存键名
            params: 参数
            file_path: 要缓存的文件路径
            content: 要缓存的内容（文本）

        Returns:
            是否成功
        """
        if not file_path and not content:
            return False

        cache_key = self._get_cache_key(key, params)
        cache_path = self._get_cache_path(cache_key)

        # 保存缓存文件
        if file_path:
            if not file_path.exists():
                return False
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(file_path, cache_path)
        elif content:
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            with open(cache_path, "w", encoding="utf-8") as f:
                f.write(content)

        # 更新元数据
        self.metadata["items"][cache_key] = {
            "key": key,
            "params": params,
            "cached_at": datetime.now().isoformat(),
            "file": str(cache_path),
        }
        self._save_metadata()

        return True

    def exists(self, key: str, params: Dict[str, Any]) -> bool:
        """
        检查缓存是否存在

        Args:
            key: 缓存键名
            params: 参数

        Returns:
            是否存在
        """
        cache_key = self._get_cache_key(key, params)
        cache_path = self._get_cache_path(cache_key)
        return cache_path.exists()

    def invalidate(self, key: Optional[str] = None) -> None:
        """
        失效缓存

        Args:
            key: 要失效的缓存键（None 表示失效所有缓存）
        """
        if key is None:
            # 清除所有缓存
            for cache_file in self.cache_dir.glob("*.cache"):
                cache_file.unlink()
            self.metadata = {
                "version": self.CACHE_VERSION,
                "items": {},
                "created_at": datetime.now().isoformat(),
            }
            self._save_metadata()
        else:
            # 清除指定键的缓存
            to_remove = [k for k in self.metadata["items"].keys() if k.startswith(key)]
            for cache_key in to_remove:
                cache_path = self._get_cache_path(cache_key)
                if cache_path.exists():
                    cache_path.unlink()
                del self.metadata["items"][cache_key]
            self._save_metadata()

    def get_stats(self) -> Dict[str, Any]:
        """获取缓存统计信息"""
        cache_files = list(self.cache_dir.glob("*.cache"))
        total_size = sum(f.stat().st_size for f in cache_files)

        return {
            "version": self.CACHE_VERSION,
            "cache_dir": str(self.cache_dir),
            "item_count": len(cache_files),
            "total_size_bytes": total_size,
            "total_size_mb": round(total_size / (1024 * 1024), 2),
            "items": list(self.metadata["items"].values()),
        }

    def cleanup(self, max_age_days: int = 30) -> int:
        """
        清理过期缓存

        Args:
            max_age_days: 最大保留天数

        Returns:
            清理的文件数
        """
        from datetime import timedelta

        cutoff_date = datetime.now() - timedelta(days=max_age_days)
        removed_count = 0

        to_remove = []
        for cache_key, item in self.metadata["items"].items():
            cached_at = item.get("cached_at")
            if cached_at:
                try:
                    cached_date = datetime.fromisoformat(cached_at)
                    if cached_date < cutoff_date:
                        to_remove.append(cache_key)
                except Exception:
                    pass

        for cache_key in to_remove:
            cache_path = self._get_cache_path(cache_key)
            if cache_path.exists():
                cache_path.unlink()
            del self.metadata["items"][cache_key]
            removed_count += 1

        if to_remove:
            self._save_metadata()

        return removed_count


class PipelineCache:
    """流水线缓存管理器

    用于管理视频生成流水线中各个阶段的缓存
    """

    def __init__(self, cache_dir: Path, enabled: bool = True):
        """
        初始化流水线缓存

        Args:
            cache_dir: 缓存目录
            enabled: 是否启用缓存
        """
        self.cache_dir = cache_dir
        self.enabled = enabled
        self.manager = CacheManager(cache_dir) if enabled else None

    def cache_scenes(
        self,
        input_file: Path,
        scenes_data: Dict[str, Any],
        use_llm: bool = False,
        provider: Optional[str] = None,
    ) -> None:
        """缓存场景数据"""
        if not self.enabled or not self.manager:
            return

        params = {
            "input_file": str(input_file),
            "use_llm": use_llm,
            "provider": provider,
        }
        content = json.dumps(scenes_data, ensure_ascii=False, indent=2)
        self.manager.set("scenes", params, content=content)

    def get_cached_scenes(
        self,
        input_file: Path,
        use_llm: bool = False,
        provider: Optional[str] = None,
    ) -> Optional[Dict[str, Any]]:
        """获取缓存的场景数据"""
        if not self.enabled or not self.manager:
            return None

        params = {
            "input_file": str(input_file),
            "use_llm": use_llm,
            "provider": provider,
        }
        cache_path = self.manager.get("scenes", params)
        if not cache_path:
            return None

        try:
            with open(cache_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            return None

    def cache_audio(
        self,
        narration: str,
        voice: str,
        rate: str,
        pitch: str,
        audio_path: Path,
    ) -> None:
        """缓存音频文件"""
        if not self.enabled or not self.manager:
            return

        params = {
            "narration_hash": self.manager._get_content_hash(narration),
            "voice": voice,
            "rate": rate,
            "pitch": pitch,
        }
        self.manager.set("audio", params, file_path=audio_path)

    def get_cached_audio(
        self,
        narration: str,
        voice: str,
        rate: str,
        pitch: str,
        output_path: Path,
    ) -> bool:
        """获取缓存的音频文件"""
        if not self.enabled or not self.manager:
            return False

        params = {
            "narration_hash": self.manager._get_content_hash(narration),
            "voice": voice,
            "rate": rate,
            "pitch": pitch,
        }
        cache_path = self.manager.get("audio", params, output_path=output_path)
        return cache_path is not None

    def cache_subtitles(
        self,
        audio_hash: str,
        subtitles_path: Path,
    ) -> None:
        """缓存字幕文件"""
        if not self.enabled or not self.manager:
            return

        params = {
            "audio_hash": audio_hash,
        }
        self.manager.set("subtitles", params, file_path=subtitles_path)

    def get_cached_subtitles(
        self,
        audio_hash: str,
        output_path: Path,
    ) -> bool:
        """获取缓存的字幕文件"""
        if not self.enabled or not self.manager:
            return False

        params = {
            "audio_hash": audio_hash,
        }
        cache_path = self.manager.get("subtitles", params, output_path=output_path)
        return cache_path is not None

    def cache_manim_script(
        self,
        scenes_hash: str,
        captions_hash: str,
        resolution: str,
        fps: int,
        provider: Optional[str],
        script_path: Path,
    ) -> None:
        """缓存 Manim 脚本"""
        if not self.enabled or not self.manager:
            return

        params = {
            "scenes_hash": scenes_hash,
            "captions_hash": captions_hash,
            "resolution": resolution,
            "fps": fps,
            "provider": provider,
        }
        self.manager.set("manim_script", params, file_path=script_path)

    def get_cached_manim_script(
        self,
        scenes_hash: str,
        captions_hash: str,
        resolution: str,
        fps: int,
        provider: Optional[str],
        output_path: Path,
    ) -> bool:
        """获取缓存的 Manim 脚本"""
        if not self.enabled or not self.manager:
            return False

        params = {
            "scenes_hash": scenes_hash,
            "captions_hash": captions_hash,
            "resolution": resolution,
            "fps": fps,
            "provider": provider,
        }
        cache_path = self.manager.get("manim_script", params, output_path=output_path)
        return cache_path is not None

    def cache_video(
        self,
        script_hash: str,
        resolution: str,
        fps: int,
        video_path: Path,
    ) -> None:
        """缓存渲染视频"""
        if not self.enabled or not self.manager:
            return

        params = {
            "script_hash": script_hash,
            "resolution": resolution,
            "fps": fps,
        }
        self.manager.set("video", params, file_path=video_path)

    def get_cached_video(
        self,
        script_hash: str,
        resolution: str,
        fps: int,
        output_path: Path,
    ) -> bool:
        """获取缓存的渲染视频"""
        if not self.enabled or not self.manager:
            return False

        params = {
            "script_hash": script_hash,
            "resolution": resolution,
            "fps": fps,
        }
        cache_path = self.manager.get("video", params, output_path=output_path)
        return cache_path is not None

    def invalidate(self) -> None:
        """失效所有缓存"""
        if self.manager:
            self.manager.invalidate()

    def get_stats(self) -> Dict[str, Any]:
        """获取缓存统计"""
        if self.manager:
            return self.manager.get_stats()
        return {"enabled": False}


if __name__ == "__main__":
    # 测试代码
    import tempfile

    with tempfile.TemporaryDirectory() as tmpdir:
        cache_dir = Path(tmpdir) / "cache"
        manager = CacheManager(cache_dir)

        # 测试设置和获取缓存
        test_content = "Hello, World!"
        manager.set("test", {"key": "value"}, content=test_content)

        retrieved = manager.get("test", {"key": "value"})
        if retrieved:
            print(f"✓ 缓存测试通过: {retrieved}")
        else:
            print("✗ 缓存测试失败")

        # 测试统计
        stats = manager.get_stats()
        print(f"\n缓存统计: {stats}")
